package com.os.opencv.java.chapter12;

import org.opencv.core.Core;
import org.opencv.core.CvType;
import org.opencv.core.Mat;
import org.opencv.ml.DTrees;
import org.opencv.ml.Ml;
import org.opencv.ml.TrainData;

public class Dtrees {

    public static void main(String[] args) {
        System.loadLibrary(Core.NATIVE_LIBRARY_NAME);

        //训练数据，每组两个数据，分别表示身高和体重
        float[] trainData = {180,81,178,75,185,87,168,65,
                173,68,160,52,158,55,170,63,163,60,168,66};
        Mat trainMat = new Mat(10,2, CvType.CV_32FC1);
        trainMat.put(0,0,trainData);

        //训练标签数据，0为男性，1为女性
        float[] label = {0,0,0,0,0,1,1,1,1,1};
        Mat labelMat = new Mat(10,1,CvType.CV_32FC1);
        labelMat.put(0,0,label);

        //测试数据
        float[] testData = {185,79,159,52};
        Mat testMat = new Mat(2,2,CvType.CV_32FC1);
        testMat.put(0,0,testData);

        //创建dtrees对象并设置参数
        DTrees dtrees = DTrees.create();
        dtrees.setMaxDepth(8);
        dtrees.setMinSampleCount(2);
        dtrees.setUseSurrogates(false);
        dtrees.setCVFolds(0);
        dtrees.setUse1SERule(false);
        dtrees.setTruncatePrunedTree(false);

        //训练决策树模型并输出训练结果
        TrainData td = TrainData.create(trainMat, Ml.ROW_SAMPLE, labelMat);
        boolean result = dtrees.train(td.getSamples(), Ml.ROW_SAMPLE, td.getResponses());

        if(result){
            System.out.println("training successed!");
        }else{
            System.out.println("training failed!");
        }
        System.out.println();

        //保存决策树模型
        dtrees.save("dtree.xml");

        //对测试数据进行预测并输出预测结果
        Mat predicts = new Mat();
        dtrees.predict(testMat, predicts, 0);
        System.out.println("predicts:\n" + predicts.dump());
        System.out.println();

        //用文字输出判断结果
        for(int i=0; i<predicts.height(); i++){
            int predict = (int)predicts.get(i, 0)[0];
            if(predict == 0){
                System.out.println("No." + i + " is male");
            }else if(predict == 1){
                System.out.println("No." + i + " is female");
            }
        }

    }
}
